#import ImageMetrics as IM
import Metric
import pylab
import Image
import numpy as np

import JT
import Calibration

def Test(method='gc',n_pose=50,dim_range=[0]):
    """
    This function will plot a metric over the given range, for one element of 
    the pose.
    Inputs:
        eval_range: an iterable range of positions at which to calculate the metric
        eval_param: the element of the pose that will be modified over the range
    output:
        val: array of values returned by the metric at each position in eval_range
            

    """
#===============================================================================
#    Data Paths and Parameters
#===============================================================================   
    fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
   
    inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
    staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
    calFileLat = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
    
    debugFolder = 'C:/Users/bryan/bryan-code/trunk/Images/debug/'
    resultsFolder = 'C:/Users/bryan/bryan-code/trunk/Results/'
    roi = [[210,187],[277,370]]
    body_center = (3.906, 2.17, -7.378)
    
    # ****Make sure angle and translation ranges are the same length
    # Iterate over a range of translations from -30 to 30 mm
    numel = 12
    ang = np.pi/numel
    
    
    trans_range = np.arange(-numel,numel+1,2)
    # Iterate over a range of angles from -pi/4 (45deg) to pi/4 in 2.864 deg increments
    angle_range = np.arange(-ang,ang+ang/numel,2*ang/numel)
        
#===============================================================================
#    Setup DRR
#===============================================================================  
    xraycam = JT.XRayCamera()
    drr = JT.DRR()
    drr.SetXRayCamera(xraycam)
    drr.SetVolumeFilename(inputVolFileName)
    drr.SetBlackOnWhiteImage()
    drr.InteractionOff()
    
    cal = Calibration.ExternalCalibration()
    cal.LoadConsolidatedCalibration(calFileLat)
    cal.LoadStaFile(staFile)
    
    xraycam.SetExternalCalibration(cal)
    
    drr._volume.SetOriginalTransform(cal._VolumeT)
    drr._volume.SetBodyCenter(body_center)
    #drr._volume.UseRayCastMapper()
    
    # Set Fixed image using path
    fixedImage = JT.FixedImage()
    fixedImage.SetFileName(fixedImageName)
    
    reg=JT.Registration()
    reg.SetFixedImage(fixedImage.GetImage(0))
    reg.SetMovingImageGenerator(drr)
    reg.SetImageMetric(method)

    #Based on the fluoro mask, i determined the following region of interest 
    # (row,column) format for numpy array:
    reg.SetRegionOfInterest(roi)
    
#===============================================================================
#    Start Metric evaluation loop
#==============================================================================
    T = JT.Testing()
    pose = T.GetStartPose(n_pose)
    dim1_range = trans_range
    tmp = np.ones((len(dim1_range),6))
    tmp *= pose
    tmp[:,0] = dim1_range
    all_poses = tmp
#    
#    all_poses = np.array([trans_range,trans_range,trans_range,\
#                   angle_range,angle_range,angle_range]).transpose()
#    
    num_evals = len(all_poses)
    val = np.zeros((num_evals,6))
    pylab.ion()
    fig = pylab.figure()
    fig.canvas.setFixedSize(800,600)
    #fig.set_size_inches(8,10)
    #fig.subplots_adjust(wspace=0.35,hspace=0.4)
    ax = [range(6)]
    
    for dim in dim_range:
        # loop through each dimension of the pose (Tx, Ty, Tz, Rx, Ry, Rz)
        # only want to vary one parameter at a time, all others are set to 0 using a mask
        pose_mask = np.ones((num_evals,6))
        pose_mask[:,dim]=1
        poses = all_poses*pose_mask
        for i,pose in enumerate( poses ):
            val[i,dim] = reg.GetValue(pose,debug=True)
#            print "    Metric: %s, Iteration: %i, Metric Value: %f" % \
#                  (reg._metricParams['method'],i,val[i,dim])
            #print "    Pose: ", pose
        
        #plot results
        # Create a subplot for each dimension
        ax = fig.add_subplot('11'+str(dim+1))
        ax.plot(all_poses[:,dim],val[:,dim])
        ax.set_title(reg._metricParams['method']+' Dim:'+str(dim))
        ax.set_xlabel('Pose')
        ax.set_ylabel('Metric Value')
        ax.set_xlim(all_poses[0,dim],all_poses[-1,dim])

    
    save_name = resultsFolder+method.upper()+'_metric_function_dim'+str(dim)+'_pose'+str(n_pose)+'.png'
    pylab.savefig(save_name)
    pylab.show()
    


if __name__ == "__main__":
    #method_list = ['nmi', 'mi', 'ncc', 'ecc', 'mse', 'ncc', 'gd', 'gc']
    method_list = ['gc']
    for method in method_list:
        Test(method,n_pose=10,dim_range=[0])
        print "Done ... ", method
        
        